package me.huangxiang.concurrent;

import java.io.IOException;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.nio.ByteBuffer;
import java.nio.channels.*;
import java.nio.channels.spi.SelectorProvider;
import java.util.*;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

public class NIOEchoServer {
    private Selector selector;
    private ExecutorService tp = Executors.newCachedThreadPool();
    public static Map<Socket, Long> time_start = new HashMap<>(10240);


    private void startServer() throws IOException {
        selector = SelectorProvider.provider().openSelector();
        ServerSocketChannel ssc = ServerSocketChannel.open();
        ssc.configureBlocking(false);

        InetSocketAddress isa = new InetSocketAddress("localhost", 8000);
        ssc.socket().bind(isa);

        SelectionKey acceptKey = ssc.register(selector, SelectionKey.OP_ACCEPT);

        for (;;){
            selector.select();
            Set<SelectionKey> readyKeys = selector.selectedKeys();
            Iterator<SelectionKey> i = readyKeys.iterator();
            long e = 0;
            while (i.hasNext()){
                SelectionKey sk = i.next();
                i.remove();

                if (sk.isAcceptable()){
                    doAccept(sk);
                } else if (sk.isValid() && sk.isReadable()){
                    if (!time_start.containsKey(((SocketChannel)sk.channel()).socket())){
                        time_start.put(((SocketChannel)sk.channel()).socket(), System.currentTimeMillis());
                    }
                    doRead(sk);
                }else if (sk.isValid() && sk.isWritable()){
                    doWrite(sk);
                    e = System.currentTimeMillis();
                    long b = time_start.remove(((SocketChannel)sk.channel()).socket());
                    System.out.println("spend: " + (e - b) + "ms");
                }
            }
        }
    }


    private void doAccept(SelectionKey sk){
        ServerSocketChannel server = (ServerSocketChannel)sk.channel();
        SocketChannel clientChannel = null;

        try {
            clientChannel = server.accept();
        } catch (IOException e) {
            e.printStackTrace();
        }
        try {
            clientChannel.configureBlocking(false);
        } catch (IOException e) {
            e.printStackTrace();
        }

        SelectionKey clientKey = null;
        try {
            clientKey = clientChannel.register(selector, SelectionKey.OP_READ);
        } catch (ClosedChannelException e) {
            e.printStackTrace();
        }
        EchoClient echoClient = new EchoClient();
        clientKey.attach(echoClient);

        InetAddress clientAddress = clientChannel.socket().getInetAddress();
        System.out.println("Accept connection from " + clientAddress.getHostAddress() + ".");
    }

    private void doRead(SelectionKey sk){
        SocketChannel channel = (SocketChannel)sk.channel();

        ByteBuffer bb = ByteBuffer.allocate(8192);

        int len;

        try {
            len = channel.read(bb);
            if (len < 0){
                channel.close();
            }
        } catch (IOException e) {
            e.printStackTrace();
        }

        bb.flip();
        tp.execute(new HandleMsg(sk, bb));
    }

    private void doWrite(SelectionKey sk) throws IOException {
        SocketChannel channel = (SocketChannel)sk.channel();
        EchoClient echoClient = (EchoClient)sk.attachment();
        LinkedList<ByteBuffer> outq = echoClient.getOutputQueue();

        ByteBuffer bb = outq.getLast();

        int len = channel.write(bb);
        if (bb.remaining() == 0){
            outq.removeLast();
        }

        if (outq.size() == 0){
            sk.interestOps(SelectionKey.OP_READ);
        }
    }


    class HandleMsg implements Runnable{
        SelectionKey sk;
        ByteBuffer bb;

        public HandleMsg(SelectionKey sk, ByteBuffer bb) {
            this.sk = sk;
            this.bb = bb;
        }

        @Override
        public void run() {
            EchoClient echoClient = (EchoClient) sk.attachment();
            echoClient.enqueue(bb);
            sk.interestOps(SelectionKey.OP_READ | SelectionKey.OP_WRITE);
            selector.wakeup();
        }
    }
}
